package com.coocaa.typhoon.api.v1.milvus.service;


import com.coocaa.typhoon.api.v1.milvus.constant.MilvusSchemaConstant;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;

/**
 * @author bijiahao
 * @date : 2025/9/25 10:13
 * @description
 */
@Slf4j
@Service
public class UserFileCollectionCreateService {
    @Autowired
    private MilvusClientV2 client;

    public void create(String collectionName) {
        if (StringUtils.isBlank(collectionName)) {
            log.error("collection name is blank");
            return;
        }
        if (Objects.isNull(client)) {
            throw new RuntimeException("Milvus client init failed!");
        }
        // check collection exist
        Boolean exist = client.hasCollection(HasCollectionReq.builder().collectionName(collectionName).build());
        if (exist) {
            log.info("collection name:{} already exist", collectionName);
            return;
        }

        // 定义 schema
        CreateCollectionReq.CollectionSchema schema = getCollectionSchema(client);
        // index definition
//        Map<String, Object> params = new HashMap<String, Object>() {{
//            put("nlist", 2048);
//        }};
        Map<String, Object> params = new HashMap<String, Object>() {{
            put("M", 16);
            put("efConstruction", 200);
        }};
        IndexParam tosVectorIndex = IndexParam.builder()
                .fieldName(MilvusSchemaConstant.UserFile.Column.TOS_VECTOR)
                .indexName(MilvusSchemaConstant.UserFile.Index.TOS_VECTOR)
//                .indexType(IndexParam.IndexType.IVF_SQ8)
                .indexType(IndexParam.IndexType.HNSW)
                .extraParams(params)
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        // //create sparse index
        IndexParam remarkSparseIndex = IndexParam.builder()
                .fieldName(MilvusSchemaConstant.UserFile.Column.REMARK_SPARSE_VECTOR)
                .indexName(MilvusSchemaConstant.UserFile.Index.REMARK_SPARSE_VECTOR)
                .indexType(IndexParam.IndexType.SPARSE_INVERTED_INDEX)
                .extraParams(params)
                .metricType(IndexParam.MetricType.IP)
                .build();

        // 创建 collection
        CreateCollectionReq createReq = CreateCollectionReq.builder()
                .collectionName(collectionName)
                .description("用户文件")
                .collectionSchema(schema)
                .numShards(2)
//                .indexParams(Arrays.asList(tosVectorIndex, remarkSparseIndex))
                .indexParams(Arrays.asList(tosVectorIndex))
                .enableDynamicField(Boolean.FALSE)
                .build();


        client.createCollection(createReq);
        log.info("创建成功");
        client.close();
    }

    private static CreateCollectionReq.CollectionSchema getCollectionSchema(MilvusClientV2 client) {
        CreateCollectionReq.CollectionSchema schema = client.createSchema();
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.ID).dataType(DataType.Int64).isPrimaryKey(true).autoID(false).build());
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.USER_ID).dataType(DataType.VarChar).maxLength(33).build());
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.FILE_TYPE).dataType(DataType.VarChar).maxLength(30).build());
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.FILE_PATH).dataType(DataType.VarChar).maxLength(100).build());

//        Map<String, Object> analyzerParams = new HashMap<>();
//        analyzerParams.put("type", "chinese");
//        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.REMARK).dataType(DataType.VarChar).maxLength(1000).enableAnalyzer(Boolean.TRUE).analyzerParams(analyzerParams).build());
//        schema.addField(AddFieldReq.builder().fieldName(MilvusParamAiCustomer.V3CustomerFaq.Param.STYLE_ANSWER).dataType(DataType.VarChar).maxLength(10000).build());

        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.SHOOT_TIME).dataType(DataType.Int64).build());
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.UPLOAD_TIME).dataType(DataType.Int64).build());
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.STATUS).dataType(DataType.Int8).build());
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.LOCATION).dataType(DataType.VarChar).maxLength(30).build());

        // 向量字段
        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.TOS_VECTOR).dataType(DataType.FloatVector).dimension(2048).build());
//        schema.addField(AddFieldReq.builder().fieldName(MilvusSchemaConstant.UserFile.Column.REMARK_SPARSE_VECTOR).dataType(DataType.SparseFloatVector).build());
        return schema;
    }
}
